// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_GPU_H
#define EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_GPU_H

namespace Eigen {
namespace internal {

#if defined(EIGEN_USE_GPU) && defined(EIGEN_GPUCC)
    // Full reducers for GPU, don't vectorize for now

    // Reducer function that enables multiple gpu thread to safely accumulate at the same
    // output address. It basically reads the current value of the output variable, and
    // attempts to update it with the new value. If in the meantime another gpu thread
    // updated the content of the output address it will try again.
    template <typename T, typename R> __device__ EIGEN_ALWAYS_INLINE void atomicReduce(T* output, T accum, R& reducer)
    {
#if (defined(EIGEN_HIP_DEVICE_COMPILE) && defined(__HIP_ARCH_HAS_WARP_SHUFFLE__)) || (EIGEN_CUDA_ARCH >= 300)
        if (sizeof(T) == 4)
        {
            unsigned int oldval = *reinterpret_cast<unsigned int*>(output);
            unsigned int newval = oldval;
            reducer.reduce(accum, reinterpret_cast<T*>(&newval));
            if (newval == oldval)
            {
                return;
            }
            unsigned int readback;
            while ((readback = atomicCAS((unsigned int*)output, oldval, newval)) != oldval)
            {
                oldval = readback;
                newval = oldval;
                reducer.reduce(accum, reinterpret_cast<T*>(&newval));
                if (newval == oldval)
                {
                    return;
                }
            }
        }
        else if (sizeof(T) == 8)
        {
            unsigned long long oldval = *reinterpret_cast<unsigned long long*>(output);
            unsigned long long newval = oldval;
            reducer.reduce(accum, reinterpret_cast<T*>(&newval));
            if (newval == oldval)
            {
                return;
            }
            unsigned long long readback;
            while ((readback = atomicCAS((unsigned long long*)output, oldval, newval)) != oldval)
            {
                oldval = readback;
                newval = oldval;
                reducer.reduce(accum, reinterpret_cast<T*>(&newval));
                if (newval == oldval)
                {
                    return;
                }
            }
        }
        else
        {
            gpu_assert(0 && "Wordsize not supported");
        }
#else   // EIGEN_CUDA_ARCH >= 300
        gpu_assert(0 && "Shouldn't be called on unsupported device");
#endif  // EIGEN_CUDA_ARCH >= 300
    }

    // We extend atomicExch to support extra data types
    template <typename Type> __device__ inline Type atomicExchCustom(Type* address, Type val) { return atomicExch(address, val); }

    template <> __device__ inline double atomicExchCustom(double* address, double val)
    {
        unsigned long long int* address_as_ull = reinterpret_cast<unsigned long long int*>(address);
        return __longlong_as_double(atomicExch(address_as_ull, __double_as_longlong(val)));
    }

#ifdef EIGEN_HAS_GPU_FP16
    template <typename R> __device__ inline void atomicReduce(half2* output, half2 accum, R& reducer)
    {
        unsigned int oldval = *reinterpret_cast<unsigned int*>(output);
        unsigned int newval = oldval;
        reducer.reducePacket(accum, reinterpret_cast<half2*>(&newval));
        if (newval == oldval)
        {
            return;
        }
        unsigned int readback;
        while ((readback = atomicCAS((unsigned int*)output, oldval, newval)) != oldval)
        {
            oldval = readback;
            newval = oldval;
            reducer.reducePacket(accum, reinterpret_cast<half2*>(&newval));
            if (newval == oldval)
            {
                return;
            }
        }
    }
    // reduction should be associative since reduction is not atomic in wide vector but atomic in half2 operations
    template <typename R> __device__ inline void atomicReduce(Packet4h2* output, Packet4h2 accum, R& reducer)
    {
        half2* houtput = reinterpret_cast<half2*>(output);
        half2* haccum = reinterpret_cast<half2*>(&accum);
        for (int i = 0; i < 4; ++i) { atomicReduce(houtput + i, *(haccum + i), reducer); }
    }
#endif  // EIGEN_HAS_GPU_FP16

    template <> __device__ inline void atomicReduce(float* output, float accum, SumReducer<float>&)
    {
#if (defined(EIGEN_HIP_DEVICE_COMPILE) && defined(__HIP_ARCH_HAS_WARP_SHUFFLE__)) || (EIGEN_CUDA_ARCH >= 300)
        atomicAdd(output, accum);
#else   // EIGEN_CUDA_ARCH >= 300
        gpu_assert(0 && "Shouldn't be called on unsupported device");
#endif  // EIGEN_CUDA_ARCH >= 300
    }

    template <typename CoeffType, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionInitKernel(const CoeffType val, Index num_preserved_coeffs, CoeffType* output)
    {
        const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
        const Index num_threads = blockDim.x * gridDim.x;
        for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) { output[i] = val; }
    }

    template <int BlockSize, int NumPerThread, typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void
    FullReductionKernel(Reducer reducer, const Self input, Index num_coeffs, typename Self::CoeffReturnType* output, unsigned int* semaphore)
    {
#if (defined(EIGEN_HIP_DEVICE_COMPILE) && defined(__HIP_ARCH_HAS_WARP_SHUFFLE__)) || (EIGEN_CUDA_ARCH >= 300)
        // Initialize the output value
        const Index first_index = blockIdx.x * BlockSize * NumPerThread + threadIdx.x;
        if (gridDim.x == 1)
        {
            if (first_index == 0)
            {
                *output = reducer.initialize();
            }
        }
        else
        {
            if (threadIdx.x == 0)
            {
                unsigned int block = atomicCAS(semaphore, 0u, 1u);
                if (block == 0)
                {
                    // We're the first block to run, initialize the output value
                    atomicExchCustom(output, reducer.initialize());
                    __threadfence();
                    atomicExch(semaphore, 2u);
                }
                else
                {
                    // Wait for the first block to initialize the output value.
                    // Use atomicCAS here to ensure that the reads aren't cached
                    unsigned int val;
                    do
                    {
                        val = atomicCAS(semaphore, 2u, 2u);
                    } while (val < 2u);
                }
            }
        }

        __syncthreads();

        eigen_assert(gridDim.x == 1 || *semaphore >= 2u);

        typename Self::CoeffReturnType accum = reducer.initialize();
        Index max_iter = numext::mini<Index>(num_coeffs - first_index, NumPerThread * BlockSize);
        for (Index i = 0; i < max_iter; i += BlockSize)
        {
            const Index index = first_index + i;
            eigen_assert(index < num_coeffs);
            typename Self::CoeffReturnType val = input.m_impl.coeff(index);
            reducer.reduce(val, &accum);
        }

#pragma unroll
        for (int offset = warpSize / 2; offset > 0; offset /= 2)
        {
#if defined(EIGEN_HIPCC)
            // use std::is_floating_point to determine the type of reduced_val
            // This is needed because when Type == double, hipcc will give a "call to __shfl_down is ambguous" error
            // and list the float and int versions of __shfl_down as the candidate functions.
            if (std::is_floating_point<typename Self::CoeffReturnType>::value)
            {
                reducer.reduce(__shfl_down(static_cast<float>(accum), offset, warpSize), &accum);
            }
            else
            {
                reducer.reduce(__shfl_down(static_cast<int>(accum), offset, warpSize), &accum);
            }
#elif defined(EIGEN_CUDA_SDK_VER) && EIGEN_CUDA_SDK_VER < 90000
            reducer.reduce(__shfl_down(accum, offset, warpSize), &accum);
#else
            reducer.reduce(__shfl_down_sync(0xFFFFFFFF, accum, offset, warpSize), &accum);
#endif
        }

        if ((threadIdx.x & (warpSize - 1)) == 0)
        {
            atomicReduce(output, accum, reducer);
        }

        if (gridDim.x > 1 && threadIdx.x == 0)
        {
            // Let the last block reset the semaphore
            atomicInc(semaphore, gridDim.x + 1);
#if defined(EIGEN_HIPCC)
            __threadfence_system();
#endif
        }
#else   // EIGEN_CUDA_ARCH >= 300
        gpu_assert(0 && "Shouldn't be called on unsupported device");
#endif  // EIGEN_CUDA_ARCH >= 300
    }

#ifdef EIGEN_HAS_GPU_FP16
    template <typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void
    ReductionInitFullReduxKernelHalfFloat(Reducer reducer, const Self input, Index num_coeffs, packet_traits<Eigen::half>::type* scratch)
    {
        eigen_assert(blockDim.x == 1);
        eigen_assert(gridDim.x == 1);
        typedef packet_traits<Eigen::half>::type packet_type;
        Index packet_remainder = num_coeffs % Index(unpacket_traits<packet_type>::size);
        if (packet_remainder != 0)
        {
            half2* h2scratch = reinterpret_cast<half2*>(scratch);
            for (Index i = num_coeffs - packet_remainder; i + 2 <= num_coeffs; i += 2)
            {
                *h2scratch = __halves2half2(input.m_impl.coeff(i), input.m_impl.coeff(i + 1));
                h2scratch++;
            }
            if ((num_coeffs & 1) != 0)
            {
                half lastCoeff = input.m_impl.coeff(num_coeffs - 1);
                *h2scratch = __halves2half2(lastCoeff, reducer.initialize());
            }
        }
        else
        {
            *scratch = reducer.template initializePacket<packet_type>();
        }
    }

    template <typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionInitKernelHalfFloat(Reducer reducer, const Self input, Index num_coeffs, half* output)
    {
        const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
        const Index num_threads = blockDim.x * gridDim.x;
        typedef typename packet_traits<Eigen::half>::type PacketType;

        const Index num_packets = num_coeffs / Index(unpacket_traits<PacketType>::size);
        PacketType* p_output = reinterpret_cast<PacketType*>(output);
        for (Index i = thread_id; i < num_packets; i += num_threads) { p_output[i] = reducer.template initializePacket<PacketType>(); }
        Index packet_remainder = num_coeffs % Index(unpacket_traits<PacketType>::size);
        if (thread_id < packet_remainder)
        {
            output[num_coeffs - packet_remainder + thread_id] = reducer.initialize();
        }
    }

    template <int BlockSize, int NumPerThread, typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void
    FullReductionKernelHalfFloat(Reducer reducer, const Self input, Index num_coeffs, half* output, packet_traits<Eigen::half>::type* scratch)
    {
        typedef typename packet_traits<Eigen::half>::type PacketType;
        const int packet_width = unpacket_traits<PacketType>::size;
        eigen_assert(NumPerThread % packet_width == 0);
        const Index first_index = blockIdx.x * BlockSize * NumPerThread + packet_width * threadIdx.x;

        // Initialize the output value if it wasn't initialized by the ReductionInitKernel

        if (gridDim.x == 1)
        {
            if (first_index == 0)
            {
                int rem = num_coeffs % packet_width;
                if (rem != 0)
                {
                    half2* p_scratch = reinterpret_cast<half2*>(scratch);
                    *scratch = reducer.template initializePacket<PacketType>();
                    for (int i = 0; i < rem / 2; i++)
                    {
                        *p_scratch =
                            __halves2half2(input.m_impl.coeff(num_coeffs - packet_width + 2 * i), input.m_impl.coeff(num_coeffs - packet_width + 2 * i + 1));
                        p_scratch++;
                    }
                    if ((num_coeffs & 1) != 0)
                    {
                        half last = input.m_impl.coeff(num_coeffs - 1);
                        *p_scratch = __halves2half2(last, reducer.initialize());
                    }
                }
                else
                {
                    *scratch = reducer.template initializePacket<PacketType>();
                }
            }
            __syncthreads();
        }

        PacketType accum = reducer.template initializePacket<PacketType>();
        const Index max_iter = numext::mini<Index>((num_coeffs - first_index) / packet_width, NumPerThread * BlockSize / packet_width);
        for (Index i = 0; i < max_iter; i += BlockSize)
        {
            const Index index = first_index + packet_width * i;
            eigen_assert(index + packet_width < num_coeffs);
            PacketType val = input.m_impl.template packet<Unaligned>(index);
            reducer.reducePacket(val, &accum);
        }

#pragma unroll
        for (int offset = warpSize / 2; offset > 0; offset /= 2)
        {
#if defined(EIGEN_HIPCC)
            PacketType r1;
            half2* hr = reinterpret_cast<half2*>(&r1);
            half2* hacc = reinterpret_cast<half2*>(&accum);
            for (int i = 0; i < packet_width / 2; i++)
            {
                // FIXME : remove this workaround once we have native half/half2 support for __shfl_down
                union
                {
                    int i;
                    half2 h;
                } wka_in, wka_out;
                wka_in.h = hacc[i];
                wka_out.i = __shfl_down(wka_in.i, offset, warpSize);
                hr[i] = wka_out.h;
            }
            reducer.reducePacket(r1, &accum);
#elif defined(EIGEN_CUDA_SDK_VER) && EIGEN_CUDA_SDK_VER < 90000
            PacketType r1;
            half2* hr = reinterpret_cast<half2*>(&r1);
            half2* hacc = reinterpret_cast<half2*>(&accum);
            for (int i = 0; i < packet_width / 2; i++) { hr[i] = __shfl_down(hacc[i], offset, warpSize); }
            reducer.reducePacket(r1, &accum);
#else
            PacketType r1;
            half2* hr = reinterpret_cast<half2*>(&r1);
            half2* hacc = reinterpret_cast<half2*>(&accum);
            for (int i = 0; i < packet_width / 2; i++) { hr[i] = __shfl_down_sync(0xFFFFFFFF, hacc[i], (unsigned)offset, warpSize); }
            reducer.reducePacket(r1, &accum);

#endif
        }

        if ((threadIdx.x & (warpSize - 1)) == 0)
        {
            atomicReduce(scratch, accum, reducer);
        }

        __syncthreads();
        half2* rv1 = reinterpret_cast<half2*>(scratch);
        if (packet_width > 2)
        {
            reducer.reducePacket(rv1[2], rv1);
            reducer.reducePacket(rv1[3], rv1 + 1);
            reducer.reducePacket(rv1[1], rv1);
        }
        if (gridDim.x == 1)
        {
            if (first_index == 0)
            {
                half tmp = __low2half(*rv1);
                reducer.reduce(__high2half(*rv1), &tmp);
                *output = tmp;
            }
        }
    }

    template <typename Op>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ReductionCleanupKernelHalfFloat(Op reducer, half* output, packet_traits<Eigen::half>::type* scratch)
    {
        eigen_assert(threadIdx.x == 1);
        half2* pscratch = reinterpret_cast<half2*>(scratch);
        half tmp = __float2half(0.f);
        typedef packet_traits<Eigen::half>::type packet_type;
        for (int i = 0; i < unpacket_traits<packet_type>::size; i += 2)
        {
            reducer.reduce(__low2half(*pscratch), &tmp);
            reducer.reduce(__high2half(*pscratch), &tmp);
            pscratch++;
        }
        *output = tmp;
    }

#endif  // EIGEN_HAS_GPU_FP16

    template <typename Self, typename Op, typename OutputType, bool PacketAccess, typename Enabled = void> struct FullReductionLauncher
    {
        static void run(const Self&, Op&, const GpuDevice&, OutputType*, typename Self::Index)
        {
            gpu_assert(false && "Should only be called on doubles, floats and half floats");
        }
    };

    // Specialization for float and double
    template <typename Self, typename Op, typename OutputType, bool PacketAccess>
    struct FullReductionLauncher<
        Self,
        Op,
        OutputType,
        PacketAccess,
        typename internal::enable_if<internal::is_same<float, OutputType>::value || internal::is_same<double, OutputType>::value, void>::type>
    {
        static void run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output, typename Self::Index num_coeffs)
        {
            typedef typename Self::Index Index;
            const int block_size = 256;
            const int num_per_thread = 128;
            const int num_blocks = divup<int>(num_coeffs, block_size * num_per_thread);

            unsigned int* semaphore = NULL;
            if (num_blocks > 1)
            {
                semaphore = device.semaphore();
            }

            LAUNCH_GPU_KERNEL((FullReductionKernel<block_size, num_per_thread, Self, Op, Index>),
                              num_blocks,
                              block_size,
                              0,
                              device,
                              reducer,
                              self,
                              num_coeffs,
                              output,
                              semaphore);
        }
    };

#ifdef EIGEN_HAS_GPU_FP16
    template <typename Self, typename Op> struct FullReductionLauncher<Self, Op, Eigen::half, false>
    {
        static void run(const Self&, Op&, const GpuDevice&, half*, typename Self::Index)
        {
            gpu_assert(false && "Should not be called since there is no packet accessor");
        }
    };

    template <typename Self, typename Op> struct FullReductionLauncher<Self, Op, Eigen::half, true>
    {
        static void run(const Self& self, Op& reducer, const GpuDevice& device, half* output, typename Self::Index num_coeffs)
        {
            typedef typename Self::Index Index;
            typedef typename packet_traits<Eigen::half>::type PacketType;

            const int block_size = 256;
            const int num_per_thread = 128;
            const int num_blocks = divup<int>(num_coeffs, block_size * num_per_thread);
            PacketType* scratch = static_cast<PacketType*>(device.scratchpad());
            // half2* scratch = static_cast<half2*>(device.scratchpad());

            if (num_blocks > 1)
            {
                // We initialize the output and the scrathpad outside the reduction kernel when we can't be sure that there
                // won't be a race conditions between multiple thread blocks.
                LAUNCH_GPU_KERNEL((ReductionInitFullReduxKernelHalfFloat<Self, Op, Index>), 1, 1, 0, device, reducer, self, num_coeffs, scratch);
            }

            LAUNCH_GPU_KERNEL((FullReductionKernelHalfFloat<block_size, num_per_thread, Self, Op, Index>),
                              num_blocks,
                              block_size,
                              0,
                              device,
                              reducer,
                              self,
                              num_coeffs,
                              output,
                              scratch);

            if (num_blocks > 1)
            {
                LAUNCH_GPU_KERNEL((ReductionCleanupKernelHalfFloat<Op>), 1, 1, 0, device, reducer, output, scratch);
            }
        }
    };
#endif  // EIGEN_HAS_GPU_FP16

    template <typename Self, typename Op, bool Vectorizable> struct FullReducer<Self, Op, GpuDevice, Vectorizable>
    {
        // Unfortunately nvidia doesn't support well exotic types such as complex,
        // so reduce the scope of the optimized version of the code to the simple cases
        // of doubles, floats and half floats
#ifdef EIGEN_HAS_GPU_FP16
        static const bool HasOptimizedImplementation =
            !Self::ReducerTraits::IsStateful &&
            (internal::is_same<typename Self::CoeffReturnType, float>::value || internal::is_same<typename Self::CoeffReturnType, double>::value ||
             (internal::is_same<typename Self::CoeffReturnType, Eigen::half>::value && reducer_traits<Op, GpuDevice>::PacketAccess));
#else   // EIGEN_HAS_GPU_FP16
        static const bool HasOptimizedImplementation = !Self::ReducerTraits::IsStateful && (internal::is_same<typename Self::CoeffReturnType, float>::value ||
                                                                                            internal::is_same<typename Self::CoeffReturnType, double>::value);
#endif  // EIGEN_HAS_GPU_FP16

        template <typename OutputType> static void run(const Self& self, Op& reducer, const GpuDevice& device, OutputType* output)
        {
            gpu_assert(HasOptimizedImplementation && "Should only be called on doubles, floats or half floats");
            const Index num_coeffs = array_prod(self.m_impl.dimensions());
            // Don't crash when we're called with an input tensor of size 0.
            if (num_coeffs == 0)
            {
                return;
            }

            FullReductionLauncher<Self, Op, OutputType, reducer_traits<Op, GpuDevice>::PacketAccess>::run(self, reducer, device, output, num_coeffs);
        }
    };

    template <int NumPerThread, typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void
    InnerReductionKernel(Reducer reducer, const Self input, Index num_coeffs_to_reduce, Index num_preserved_coeffs, typename Self::CoeffReturnType* output)
    {
#if (defined(EIGEN_HIP_DEVICE_COMPILE) && defined(__HIP_ARCH_HAS_WARP_SHUFFLE__)) || (EIGEN_CUDA_ARCH >= 300)
        typedef typename Self::CoeffReturnType Type;
        eigen_assert(blockDim.y == 1);
        eigen_assert(blockDim.z == 1);
        eigen_assert(gridDim.y == 1);
        eigen_assert(gridDim.z == 1);

        const int unroll_times = 16;
        eigen_assert(NumPerThread % unroll_times == 0);

        const Index input_col_blocks = divup<Index>(num_coeffs_to_reduce, blockDim.x * NumPerThread);
        const Index num_input_blocks = input_col_blocks * num_preserved_coeffs;

        const Index num_threads = blockDim.x * gridDim.x;
        const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;

        // Initialize the output values if they weren't initialized by the ReductionInitKernel
        if (gridDim.x == 1)
        {
            for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) { output[i] = reducer.initialize(); }
            __syncthreads();
        }

        for (Index i = blockIdx.x; i < num_input_blocks; i += gridDim.x)
        {
            const Index row = i / input_col_blocks;

            if (row < num_preserved_coeffs)
            {
                const Index col_block = i % input_col_blocks;
                const Index col_begin = col_block * blockDim.x * NumPerThread + threadIdx.x;

                Type reduced_val = reducer.initialize();

                for (Index j = 0; j < NumPerThread; j += unroll_times)
                {
                    const Index last_col = col_begin + blockDim.x * (j + unroll_times - 1);
                    if (last_col >= num_coeffs_to_reduce)
                    {
                        for (Index col = col_begin + blockDim.x * j; col < num_coeffs_to_reduce; col += blockDim.x)
                        {
                            const Type val = input.m_impl.coeff(row * num_coeffs_to_reduce + col);
                            reducer.reduce(val, &reduced_val);
                        }
                        break;
                    }
                    else
                    {
                        // Faster version of the loop with no branches after unrolling.
#pragma unroll
                        for (int k = 0; k < unroll_times; ++k)
                        {
                            const Index col = col_begin + blockDim.x * (j + k);
                            reducer.reduce(input.m_impl.coeff(row * num_coeffs_to_reduce + col), &reduced_val);
                        }
                    }
                }

#pragma unroll
                for (int offset = warpSize / 2; offset > 0; offset /= 2)
                {
#if defined(EIGEN_HIPCC)
                    // use std::is_floating_point to determine the type of reduced_val
                    // This is needed because when Type == double, hipcc will give a "call to __shfl_down is ambguous" error
                    // and list the float and int versions of __shfl_down as the candidate functions.
                    if (std::is_floating_point<Type>::value)
                    {
                        reducer.reduce(__shfl_down(static_cast<float>(reduced_val), offset), &reduced_val);
                    }
                    else
                    {
                        reducer.reduce(__shfl_down(static_cast<int>(reduced_val), offset), &reduced_val);
                    }
#elif defined(EIGEN_CUDA_SDK_VER) && EIGEN_CUDA_SDK_VER < 90000
                    reducer.reduce(__shfl_down(reduced_val, offset), &reduced_val);
#else
                    reducer.reduce(__shfl_down_sync(0xFFFFFFFF, reduced_val, offset), &reduced_val);
#endif
                }

                if ((threadIdx.x & (warpSize - 1)) == 0)
                {
                    atomicReduce(&(output[row]), reduced_val, reducer);
                }
            }
        }
#else   // EIGEN_CUDA_ARCH >= 300
        gpu_assert(0 && "Shouldn't be called on unsupported device");
#endif  // EIGEN_CUDA_ARCH >= 300
    }

#ifdef EIGEN_HAS_GPU_FP16

    template <int NumPerThread, typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void
    InnerReductionKernelHalfFloat(Reducer reducer, const Self input, Index num_coeffs_to_reduce, Index num_preserved_coeffs, half* output)
    {
        eigen_assert(blockDim.y == 1);
        eigen_assert(blockDim.z == 1);
        eigen_assert(gridDim.y == 1);
        eigen_assert(gridDim.z == 1);

        typedef typename packet_traits<Eigen::half>::type PacketType;
        const int packet_width = unpacket_traits<PacketType>::size;
        const int unroll_times = 16 / packet_width;
        eigen_assert(NumPerThread % unroll_times == 0);
        eigen_assert(unroll_times % 2 == 0);

        const Index input_col_blocks = divup<Index>(num_coeffs_to_reduce, blockDim.x * NumPerThread * 2);
        const Index num_input_blocks = divup<Index>(input_col_blocks * num_preserved_coeffs, 2);

        const Index num_threads = blockDim.x * gridDim.x;
        const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;

        // Initialize the output values if they weren't initialized by the ReductionInitKernel
        if (gridDim.x == 1)
        {
            Index i = packet_width * thread_id;
            for (; i + packet_width <= num_preserved_coeffs; i += packet_width * num_threads)
            {
                PacketType* poutput = reinterpret_cast<PacketType*>(output + i);
                *poutput = reducer.template initializePacket<PacketType>();
            }
            if (i < num_preserved_coeffs)
            {
                output[i] = reducer.initialize();
            }
            __syncthreads();
        }

        for (Index i = blockIdx.x; i < num_input_blocks; i += gridDim.x)
        {
            const Index row = 2 * (i / input_col_blocks);  // everybody takes 2 rows

            if (row + 1 < num_preserved_coeffs)
            {
                const Index col_block = i % input_col_blocks;
                const Index col_begin = packet_width * (col_block * blockDim.x * NumPerThread + threadIdx.x);

                PacketType reduced_val1 = reducer.template initializePacket<PacketType>();
                PacketType reduced_val2 = reducer.template initializePacket<PacketType>();

                for (Index j = 0; j < NumPerThread; j += unroll_times)
                {
                    const Index last_col = col_begin + blockDim.x * (j + unroll_times - 1) * packet_width;
                    if (last_col >= num_coeffs_to_reduce)
                    {
                        Index col = col_begin + blockDim.x * j;
                        for (; col + packet_width <= num_coeffs_to_reduce; col += blockDim.x)
                        {
                            const PacketType val1 = input.m_impl.template packet<Unaligned>(row * num_coeffs_to_reduce + col);
                            reducer.reducePacket(val1, &reduced_val1);
                            const PacketType val2 = input.m_impl.template packet<Unaligned>((row + 1) * num_coeffs_to_reduce + col);
                            reducer.reducePacket(val2, &reduced_val2);
                        }
                        if (col < num_coeffs_to_reduce)
                        {
                            PacketType r1 = reducer.template initializePacket<PacketType>();
                            PacketType r2 = reducer.template initializePacket<PacketType>();
                            half2* hr1 = reinterpret_cast<half2*>(&r1);
                            half2* hr2 = reinterpret_cast<half2*>(&r2);
                            while (col + 1 < num_coeffs_to_reduce)
                            {
                                *hr1 = __halves2half2(input.m_impl.coeff(row * num_coeffs_to_reduce + col),
                                                      input.m_impl.coeff(row * num_coeffs_to_reduce + col + 1));
                                *hr2 = __halves2half2(input.m_impl.coeff((row + 1) * num_coeffs_to_reduce + col),
                                                      input.m_impl.coeff((row + 1) * num_coeffs_to_reduce + col + 1));
                                hr1++;
                                hr2++;
                                col += 2;
                            }
                            if (col < num_coeffs_to_reduce)
                            {
                                // Peel;
                                const half last1 = input.m_impl.coeff(row * num_coeffs_to_reduce + col);
                                *hr1 = __halves2half2(last1, reducer.initialize());
                                const half last2 = input.m_impl.coeff((row + 1) * num_coeffs_to_reduce + col);
                                *hr2 = __halves2half2(last2, reducer.initialize());
                            }
                            reducer.reducePacket(r1, &reduced_val1);
                            reducer.reducePacket(r2, &reduced_val2);
                        }
                        break;
                    }
                    else
                    {
                        // Faster version of the loop with no branches after unrolling.
#pragma unroll
                        for (int k = 0; k < unroll_times; ++k)
                        {
                            const Index col = col_begin + blockDim.x * (j + k) * packet_width;
                            reducer.reducePacket(input.m_impl.template packet<Unaligned>(row * num_coeffs_to_reduce + col), &reduced_val1);
                            reducer.reducePacket(input.m_impl.template packet<Unaligned>((row + 1) * num_coeffs_to_reduce + col), &reduced_val2);
                        }
                    }
                }

#pragma unroll
                for (int offset = warpSize / 2; offset > 0; offset /= 2)
                {
#if defined(EIGEN_HIPCC)
                    PacketType r1;
                    PacketType r2;
                    half2* hr1 = reinterpret_cast<half2*>(&r1);
                    half2* hr2 = reinterpret_cast<half2*>(&r2);
                    half2* rv1 = reinterpret_cast<half2*>(&reduced_val1);
                    half2* rv2 = reinterpret_cast<half2*>(&reduced_val2);
                    for (int i = 0; i < packet_width / 2; i++)
                    {
                        // FIXME : remove this workaround once we have native half/half2 support for __shfl_down
                        union
                        {
                            int i;
                            half2 h;
                        } wka_in1, wka_out1;
                        wka_in1.h = rv1[i];
                        wka_out1.i = __shfl_down(wka_in1.i, offset, warpSize);
                        hr1[i] = wka_out1.h;

                        union
                        {
                            int i;
                            half2 h;
                        } wka_in2, wka_out2;
                        wka_in2.h = rv2[i];
                        wka_out2.i = __shfl_down(wka_in2.i, offset, warpSize);
                        hr2[i] = wka_out2.h;
                    }
                    reducer.reducePacket(r1, &reduced_val1);
                    reducer.reducePacket(r2, &reduced_val2);
#elif defined(EIGEN_CUDA_SDK_VER) && EIGEN_CUDA_SDK_VER < 90000
                    PacketType r1;
                    PacketType r2;
                    half2* hr1 = reinterpret_cast<half2*>(&r1);
                    half2* hr2 = reinterpret_cast<half2*>(&r2);
                    half2* rv1 = reinterpret_cast<half2*>(&reduced_val1);
                    half2* rv2 = reinterpret_cast<half2*>(&reduced_val2);
                    for (int i = 0; i < packet_width / 2; i++)
                    {
                        hr1[i] = __shfl_down(rv1[i], offset, warpSize);
                        hr2[i] = __shfl_down(rv2[i], offset, warpSize);
                    }
                    reducer.reducePacket(r1, &reduced_val1);
                    reducer.reducePacket(r2, &reduced_val2);
#else
                    PacketType r1;
                    PacketType r2;
                    half2* hr1 = reinterpret_cast<half2*>(&r1);
                    half2* hr2 = reinterpret_cast<half2*>(&r2);
                    half2* rr1 = reinterpret_cast<half2*>(&reduced_val1);
                    half2* rr2 = reinterpret_cast<half2*>(&reduced_val2);
                    for (int i = 0; i < packet_width / 2; i++)
                    {
                        hr1[i] = __shfl_down_sync(0xFFFFFFFF, rr1[i], (unsigned)offset, warpSize);
                        hr2[i] = __shfl_down_sync(0xFFFFFFFF, rr2[i], (unsigned)offset, warpSize);
                    }
                    reducer.reducePacket(r1, &reduced_val1);
                    reducer.reducePacket(r2, &reduced_val2);

#endif
                }
                half2* rv1 = reinterpret_cast<half2*>(&reduced_val1);
                half2* rv2 = reinterpret_cast<half2*>(&reduced_val2);
                half2 val;
                if (packet_width > 2)
                {
                    reducer.reducePacket(rv1[2], rv1);
                    reducer.reducePacket(rv1[3], rv1 + 1);
                    reducer.reducePacket(rv1[1], rv1);
                    reducer.reducePacket(rv2[2], rv2);
                    reducer.reducePacket(rv2[3], rv2 + 1);
                    reducer.reducePacket(rv2[1], rv2);
                }
                half val1 = __low2half(*rv1);
                reducer.reduce(__high2half(*rv1), &val1);
                half val2 = __low2half(*rv2);
                reducer.reduce(__high2half(*rv2), &val2);
                val = __halves2half2(val1, val2);
                if ((threadIdx.x & (warpSize - 1)) == 0)
                {
                    half* loc = output + row;
                    atomicReduce((half2*)loc, val, reducer);
                }
            }
        }
    }

#endif  // EIGEN_HAS_GPU_FP16

    template <typename Self, typename Op, typename OutputType, bool PacketAccess, typename Enabled = void> struct InnerReductionLauncher
    {
        static EIGEN_DEVICE_FUNC bool run(const Self&, Op&, const GpuDevice&, OutputType*, typename Self::Index, typename Self::Index)
        {
            gpu_assert(false && "Should only be called to reduce doubles, floats and half floats on a gpu device");
            return true;
        }
    };

    // Specialization for float and double
    template <typename Self, typename Op, typename OutputType, bool PacketAccess>
    struct InnerReductionLauncher<
        Self,
        Op,
        OutputType,
        PacketAccess,
        typename internal::enable_if<internal::is_same<float, OutputType>::value || internal::is_same<double, OutputType>::value, void>::type>
    {
        static bool run(const Self& self,
                        Op& reducer,
                        const GpuDevice& device,
                        OutputType* output,
                        typename Self::Index num_coeffs_to_reduce,
                        typename Self::Index num_preserved_vals)
        {
            typedef typename Self::Index Index;

            const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals;
            const int block_size = 256;
            const int num_per_thread = 128;
            const int dyn_blocks = divup<int>(num_coeffs, block_size * num_per_thread);
            const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / block_size;
            const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);

            if (num_blocks > 1)
            {
                // We initialize the outputs outside the reduction kernel when we can't be sure that there
                // won't be a race conditions between multiple thread blocks.
                const int dyn_blocks = divup<int>(num_preserved_vals, 1024);
                const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / 1024;
                const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);
                LAUNCH_GPU_KERNEL((ReductionInitKernel<OutputType, Index>), num_blocks, 1024, 0, device, reducer.initialize(), num_preserved_vals, output);
            }

            LAUNCH_GPU_KERNEL((InnerReductionKernel<num_per_thread, Self, Op, Index>),
                              num_blocks,
                              block_size,
                              0,
                              device,
                              reducer,
                              self,
                              num_coeffs_to_reduce,
                              num_preserved_vals,
                              output);

            return false;
        }
    };

#ifdef EIGEN_HAS_GPU_FP16
    template <typename Self, typename Op> struct InnerReductionLauncher<Self, Op, Eigen::half, false>
    {
        static bool run(const Self&, Op&, const GpuDevice&, half*, typename Self::Index, typename Self::Index)
        {
            gpu_assert(false && "Should not be called since there is no packet accessor");
            return true;
        }
    };

    template <typename Self, typename Op> struct InnerReductionLauncher<Self, Op, Eigen::half, true>
    {
        static bool run(const Self& self,
                        Op& reducer,
                        const GpuDevice& device,
                        half* output,
                        typename Self::Index num_coeffs_to_reduce,
                        typename Self::Index num_preserved_vals)
        {
            typedef typename Self::Index Index;

            if (num_preserved_vals % 2 != 0)
            {
                // Not supported yet, revert to the slower code path
                return true;
            }

            const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals;
            const int block_size = /*256*/ 128;
            const int num_per_thread = /*128*/ 64;
            const int dyn_blocks = divup<int>(num_coeffs, block_size * num_per_thread);
            const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / block_size;
            const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);

            if (num_blocks > 1)
            {
                // We initialize the outputs outside the reduction kernel when we can't be sure that there
                // won't be a race conditions between multiple thread blocks.
                LAUNCH_GPU_KERNEL((ReductionInitKernelHalfFloat<Self, Op, Index>), 1, 1, 0, device, reducer, self, num_preserved_vals, output);
            }

            LAUNCH_GPU_KERNEL((InnerReductionKernelHalfFloat<num_per_thread, Self, Op, Index>),
                              num_blocks,
                              block_size,
                              0,
                              device,
                              reducer,
                              self,
                              num_coeffs_to_reduce,
                              num_preserved_vals,
                              output);

            return false;
        }
    };
#endif  // EIGEN_HAS_GPU_FP16

    template <typename Self, typename Op> struct InnerReducer<Self, Op, GpuDevice>
    {
        // Unfortunately nvidia doesn't support well exotic types such as complex,
        // so reduce the scope of the optimized version of the code to the simple case
        // of floats and half floats.
#ifdef EIGEN_HAS_GPU_FP16
        static const bool HasOptimizedImplementation =
            !Self::ReducerTraits::IsStateful &&
            (internal::is_same<typename Self::CoeffReturnType, float>::value || internal::is_same<typename Self::CoeffReturnType, double>::value ||
             (internal::is_same<typename Self::CoeffReturnType, Eigen::half>::value && reducer_traits<Op, GpuDevice>::PacketAccess));
#else   // EIGEN_HAS_GPU_FP16
        static const bool HasOptimizedImplementation = !Self::ReducerTraits::IsStateful && (internal::is_same<typename Self::CoeffReturnType, float>::value ||
                                                                                            internal::is_same<typename Self::CoeffReturnType, double>::value);
#endif  // EIGEN_HAS_GPU_FP16

        template <typename OutputType>
        static bool run(const Self& self,
                        Op& reducer,
                        const GpuDevice& device,
                        OutputType* output,
                        typename Self::Index num_coeffs_to_reduce,
                        typename Self::Index num_preserved_vals)
        {
            gpu_assert(HasOptimizedImplementation && "Should only be called on doubles, floats or half floats");
            const Index num_coeffs = array_prod(self.m_impl.dimensions());
            // Don't crash when we're called with an input tensor of size 0.
            if (num_coeffs == 0)
            {
                return true;
            }
            // It's faster to use the usual code.
            if (num_coeffs_to_reduce <= 128)
            {
                return true;
            }

            return InnerReductionLauncher<Self, Op, OutputType, reducer_traits<Op, GpuDevice>::PacketAccess>::run(
                self, reducer, device, output, num_coeffs_to_reduce, num_preserved_vals);
        }
    };

    template <int NumPerThread, typename Self, typename Reducer, typename Index>
    __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void
    OuterReductionKernel(Reducer reducer, const Self input, Index num_coeffs_to_reduce, Index num_preserved_coeffs, typename Self::CoeffReturnType* output)
    {
        const Index num_threads = blockDim.x * gridDim.x;
        const Index thread_id = blockIdx.x * blockDim.x + threadIdx.x;
        // Initialize the output values if they weren't initialized by the ReductionInitKernel
        if (gridDim.x == 1)
        {
            for (Index i = thread_id; i < num_preserved_coeffs; i += num_threads) { output[i] = reducer.initialize(); }
            __syncthreads();
        }

        // Do the reduction.
        const Index max_iter = num_preserved_coeffs * divup<Index>(num_coeffs_to_reduce, NumPerThread);
        for (Index i = thread_id; i < max_iter; i += num_threads)
        {
            const Index input_col = i % num_preserved_coeffs;
            const Index input_row = (i / num_preserved_coeffs) * NumPerThread;
            typename Self::CoeffReturnType reduced_val = reducer.initialize();
            const Index max_row = numext::mini(input_row + NumPerThread, num_coeffs_to_reduce);
            for (Index j = input_row; j < max_row; j++)
            {
                typename Self::CoeffReturnType val = input.m_impl.coeff(j * num_preserved_coeffs + input_col);
                reducer.reduce(val, &reduced_val);
            }
            atomicReduce(&(output[input_col]), reduced_val, reducer);
        }
    }

    template <typename Self, typename Op> struct OuterReducer<Self, Op, GpuDevice>
    {
        // Unfortunately nvidia doesn't support well exotic types such as complex,
        // so reduce the scope of the optimized version of the code to the simple case
        // of floats.
        static const bool HasOptimizedImplementation = !Self::ReducerTraits::IsStateful && (internal::is_same<typename Self::CoeffReturnType, float>::value ||
                                                                                            internal::is_same<typename Self::CoeffReturnType, double>::value);
        template <typename Device, typename OutputType>
        static
#if !defined(EIGEN_HIPCC)
            // FIXME :  leaving this EIGEN_DEVICE_FUNC in, results in the following runtime error
            //          (in the cxx11_tensor_reduction_gpu test)
            //
            // terminate called after throwing an instance of 'std::runtime_error'
            //   what():  No device code available for function: _ZN5Eigen8internal20OuterReductionKernelIL...
            //
            // don't know why this happens (and why is it a runtime error instead of a compile time error)
            //
            // this will be fixed by HIP PR#457
            EIGEN_DEVICE_FUNC
#endif
            bool
            run(const Self&, Op&, const Device&, OutputType*, typename Self::Index, typename Self::Index)
        {
            gpu_assert(false && "Should only be called to reduce doubles or floats on a gpu device");
            return true;
        }

        static bool run(const Self& self,
                        Op& reducer,
                        const GpuDevice& device,
                        float* output,
                        typename Self::Index num_coeffs_to_reduce,
                        typename Self::Index num_preserved_vals)
        {
            typedef typename Self::Index Index;

            // It's faster to use the usual code.
            if (num_coeffs_to_reduce <= 32)
            {
                return true;
            }

            const Index num_coeffs = num_coeffs_to_reduce * num_preserved_vals;
            const int block_size = 256;
            const int num_per_thread = 16;
            const int dyn_blocks = divup<int>(num_coeffs, block_size * num_per_thread);
            const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / block_size;
            const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);

            if (num_blocks > 1)
            {
                // We initialize the outputs in the reduction kernel itself when we don't have to worry
                // about race conditions between multiple thread blocks.
                const int dyn_blocks = divup<int>(num_preserved_vals, 1024);
                const int max_blocks = device.getNumGpuMultiProcessors() * device.maxGpuThreadsPerMultiProcessor() / 1024;
                const int num_blocks = numext::mini<int>(max_blocks, dyn_blocks);
                LAUNCH_GPU_KERNEL((ReductionInitKernel<float, Index>), num_blocks, 1024, 0, device, reducer.initialize(), num_preserved_vals, output);
            }

            LAUNCH_GPU_KERNEL((OuterReductionKernel<num_per_thread, Self, Op, Index>),
                              num_blocks,
                              block_size,
                              0,
                              device,
                              reducer,
                              self,
                              num_coeffs_to_reduce,
                              num_preserved_vals,
                              output);

            return false;
        }
    };

#endif  // defined(EIGEN_USE_GPU) && defined(EIGEN_GPUCC)

}  // end namespace internal
}  // end namespace Eigen

#endif  // EIGEN_CXX11_TENSOR_TENSOR_REDUCTION_GPU_H
